from data_process import *
from models import *
import matplotlib.pyplot as plt
import time
from config import *
import requests
import torch.nn.functional as F


"""----------------训练过程是否改变参数 ,如果不改变，把这三行注释掉，并把下面的variables['lr']改成变量lr, variables['epochs']改成变量epochs 等--------------"""
thread =threading.Thread(target=change_param) #  启动线程，用于改变模型参数
thread.daemon = True
thread.start()



device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用{device}进行训练")

network = CNN(class_num).to(device)


# 如果有已存在的模型，则加载模型，并在其基础上继续训练
if os.path.exists("./network.pth"):
    network.load_state_dict(torch.load("./network.pth"))
    print("加载已存在的模型,继续训练")

optimizer = torch.optim.Adam(network.parameters(), lr=lr)

train_loss=[]
val_precision=[]
network.train()

epoch=1
while epoch <= variables['epochs']:
    start_time = time.time()
    epoch_loss=0
    for param_group in optimizer.param_groups:  # 更新参数的学习率
        param_group['lr'] = variables['lr']
    # print("Epoch:{}  lr:{}".format(epoch, optimizer.state_dict()['param_groups'][0]['lr']))

    for images,lable in trainloader:
        images = images.to(device)
        target = lable.to(device)
        output = network(images)
        loss = F.cross_entropy(output, target)
        epoch_loss+=loss.item()
        network.zero_grad()
        loss.backward()
        optimizer.step()

    train_loss.append(epoch_loss)

    print(f"({epoch}/{variables['epochs']}),训练损失值：{epoch_loss:.2f}" ,end=' ,')
    end_time=time.time()
    print(f"本轮训练耗时：{int(end_time-start_time)}秒 ,时间：{time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))}")


    """-------------------保存模型----------------------------"""
    if int(variables['save_current_model'])==1:
        torch.save(network.state_dict(), f"./network{epoch}.pth")
        variables['save_current_model']=0

    if  epoch % save_step == 0:
        # 训练过程关注损失值的变化，一旦发现损失值上升，及时停止训练，检查代码。 发现损失值收敛，也要提前中断训练
        # 所以，要每隔几次把模型参数保存一下，防止训练过程中断，导致模型无法保存
        torch.save(network.state_dict(),f"./network{epoch}.pth")


    """----------------------模型验证----------------------------------"""
    if epoch % 2 == 0:   # 每隔两次epoch验证一次模型
        network.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, lable in valloader:
                images = images.to(device)
                lable = lable.to(device)
                outputs = network(images)
                _, predicted = torch.max(outputs.data, 1)
                total += lable.size(0)
                correct += (predicted == lable).sum().item()
            precision=100 * correct / total
        val_precision.append(precision)
        print(f"模型在第{epoch}验证集的正确率为：{precision}%")

    epoch += 1




"""由于训练的时间要分钟一轮，不能一直看着，所以训练完后发送消息到微信，提醒自己"""
# 微信消息推送用法：
# content = "模型已经训练完毕，GPU要冒烟了！！！"        # 要发给微信服务号的消息内容
# url1 = (f'https://wxpusher.zjiecode.com/api/send/message/?appToken= &content={content}&uid= &url=http%3a%2f%2fwxpusher.zjiecode.com')
# response = requests.get(url1).json()
# if response['msg'] == '处理成功':
#     print('微信消息推送成功，请查收！')
# else:
#     print('微信消息推送失败！请查看原因')


network.eval()
correct = 0
total = 0
with torch.no_grad():
        for images, lable in testloader:
                images = images.to(device)
                lable = lable.to(device)
                outputs = network(images)
                _, predicted = torch.max(outputs.data, 1)
                total += lable.size(0)
                correct += (predicted == lable).sum().item()
        precision = 100 * correct / total
print(f"模型在测试集的正确率为：{precision}%")



plt.plot(range(1,variables['epochs']+1),train_loss)
plt.show()





